# python -m lilab.OpenLabCluster_train.b1_chi_dbi clippredfile --PCA 12
from sklearn.metrics import silhouette_score, calinski_harabasz_score, davies_bouldin_score
import pickle
import argparse


clippredfile = '/mnt/liying.cibr.ac.cn_Data_Temp/multiview_9/chenxf/00_BehaviorAnalysis-seq2seq/SexAge/Day55_Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/FWPCA0.00_P100_en3_hid30_epoch264-decSeqPC0.9_svm2allAcc0.94_kmeansK2use-42_fromK1-20_K100.clippredpkl'

def main(clippredfile, pca):
    clippreddata = pickle.load(open(clippredfile,'rb'))
    cluster_labels = clippreddata['cluster_labels'] - 1 #labels start from 0
    embedding = clippreddata['embedding']
    if pca is not None and pca<embedding.shape[1]:
        embedding = embedding[:,:pca]

    dbi = davies_bouldin_score(embedding, cluster_labels)
    chi = calinski_harabasz_score(embedding, cluster_labels)
    shi = silhouette_score(embedding, cluster_labels, metric = 'correlation')
    print('dbi:', dbi) #the lower the better clustering
    print('chi:', chi) #the higher the better
    print('shi:', shi) #the smaller the better

    clippreddata['cluster_metric'] = {'davies_bouldin_score':dbi, 'calinski_harabasz_score':chi, 'silhouette_score':shi}
    pickle.dump(clippreddata, open(clippredfile,'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('clippredfile', type=str, help='clippredfile')
    parser.add_argument('--PCA', type=int, default=None)
    args = parser.parse_args()
    main(args.clippredfile, args.PCA)
